package com.suppor.rpc.register;

import com.google.common.collect.Lists;
import com.suppor.rpc.core.register.RegisterInfo;
import com.suppor.rpc.core.util.RpcServiceUtil;
import com.suppor.rpc.register.error.RegisterException;
import com.suppor.rpc.register.loadbalance.ConsistentHashLoadBalance;
import com.suppor.rpc.register.loadbalance.ServiceBalancer;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;

public class ZKRegisterCenter implements RegisterService {
    Logger LOG = LoggerFactory.getLogger(ZKRegisterCenter.class);
    /**
     * 链接失败重试间隔时间
     */
    private static final Integer baseSleepTimeMs = 1000;
    /**
     * 链接失败重试次数
     */
    private static final Integer MAX_RETRY_TIMES = 5;
    /**
     * RPC root path
     */
    private static final String SUPPOR_RPC_ROOT = "/rpc";
    private String registryAddr;
    private ServiceDiscovery<RegisterInfo> discovery;
    private ServiceBalancer<RegisterInfo> balancer;

    public ZKRegisterCenter(String registryAddr) {
        this.registryAddr = registryAddr;
        balancer = new ConsistentHashLoadBalance();
        // 1.连接Zookeeper注册中心
        CuratorFramework client = CuratorFrameworkFactory.newClient(this.registryAddr,
                new ExponentialBackoffRetry(baseSleepTimeMs, MAX_RETRY_TIMES));
        client.start();

        JsonInstanceSerializer<RegisterInfo> serializer = new JsonInstanceSerializer<>(RegisterInfo.class);

        // 2.基于Curator框架构建一个服务注册工具类
        this.discovery = ServiceDiscoveryBuilder
                .builder(RegisterInfo.class) // 自定义数据
                .client(client)
                .serializer(serializer)
                .basePath(SUPPOR_RPC_ROOT)
                .build();
        try {
            this.discovery.start();
            LOG.info("Register Center Start...");
        } catch (Exception e) {
            LOG.error("Register Center Start Error", e);
        }
    }

    @Override
    public boolean register(RegisterInfo info) throws Exception {
        String serviceKey = RpcServiceUtil.serviceKey(info.getServiceName(), info.getVersion());

        ServiceInstance<RegisterInfo> instance = ServiceInstance.<RegisterInfo>builder()
                .name(serviceKey)
                .port(info.getPort())
                .address(info.getServerAddr())
                .payload(info)
                .build();

        this.discovery.registerService(instance);
        return true;
    }

    @Override
    public boolean unRegister(RegisterInfo info) throws Exception {
        String serviceKey = RpcServiceUtil.serviceKey(info.getServiceName(), info.getVersion());
        ServiceInstance<RegisterInfo> instance = ServiceInstance.<RegisterInfo>builder()
                .name(serviceKey)
                .port(info.getPort())
                .address(info.getServerAddr())
                .payload(info)
                .build();

        this.discovery.unregisterService(instance);
        return true;
    }

    @Override
    public RegisterInfo discovery(String serviceName, String version, int hashCode) throws Exception {
        String key = RpcServiceUtil.serviceKey(serviceName, version);
        Collection<ServiceInstance<RegisterInfo>> instances = this.discovery.queryForInstances(key);

        List<RegisterInfo> services = instances.stream().map(new Function<ServiceInstance<RegisterInfo>, RegisterInfo>() {
            @Override
            public RegisterInfo apply(ServiceInstance<RegisterInfo> registerInfoServiceInstance) {
                return registerInfoServiceInstance.getPayload();
            }
        }).collect(Collectors.toList());
        RegisterInfo node = this.balancer.select(services, hashCode);
        if (null == node) {
            throw new RegisterException("Register Center Not Found Service " + serviceName + "#" + version);
        }
        return node;
    }

    @Override
    public boolean destroy() throws Exception {
        this.discovery.close();
        return true;
    }
}
